import numpy as np
import matplotlib.pyplot as plt
import sys
import os

def read_data(file_path):
    """Reads data from a file. Assumes two columns of numbers separated by whitespace."""
    data = np.loadtxt(file_path)
    return data[:, 0], data[:, 1:]  # 第一列是时间，后面所有列是状态变量

def plot_data(x, y, output_name, title="Plot", xlabel="X-axis", ylabel="Y-axis"):
    """Plots the data."""
    plt.figure(figsize=(10, 8))
    
    # 检查y是否是多列数据
    if len(y.shape) > 1 and y.shape[1] > 1:
        # 如果是多列，则绘制每一列
        for i in range(y.shape[1]):
            plt.plot(x, y[:, i], label=f'u_{i+1}')
    else:
        # 如果是单列，直接绘制
        plt.plot(x, y, marker='o', markersize=1, linestyle='-', color='b')
    
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True)
    # plt.legend()
    
    # 保存图像
    plt.savefig(output_name, dpi=300, bbox_inches='tight')
    print(f"plot save as {output_name}")

def plot_3d_trajectory(x, y, output_name):
    """创建3D轨道图"""
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 假设y的前三列是坐标
    if y.shape[1] >= 3:
        ax.plot(y[:, 0], y[:, 1], y[:, 2], marker='o', markersize=1)
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('3-body trajectory')
        
        # 保存3D图
        plt.savefig(output_name, dpi=300, bbox_inches='tight')
        print(f"3D trajectoy saved as {output_name}")
    else:
        print("data does not have enough columns for 3D plot.")

if __name__ == "__main__":
    # 处理命令行参数
    if len(sys.argv) > 1:
        file_path = sys.argv[1]
    else:
        file_path = "output.txt"  # 默认文件名
    
    try:
        # 读取数据
        x, y = read_data(file_path)
        
        # 生成输出文件名（保持文件名，更改扩展名为png）
        base_name = os.path.splitext(os.path.basename(file_path))[0]
        output_file = f"{base_name}.png"
        
        # 绘制常规图
        plot_data(x, y, output_file, 
                 title=f"solution trace - {base_name}", 
                 xlabel="t", 
                 ylabel="u(t)")
        
        # 如果数据有至少3列，也绘制3D轨道图
        if y.shape[1] >= 3:
            output_3d_file = f"{base_name}_3d_trajectory.png"
            plot_3d_trajectory(x, y, output_3d_file)
            
    except Exception as e:
        print(f"error: {e}")